import argparse

import torch

from model_trainer import BYOLTrainer
from model import ResNetBYOL, MLPHead

torch.manual_seed(0)

parser = argparse.ArgumentParser(description='Self-Supervised Learning - BYOL - PyTorch')
parser.add_argument('-data', metavar='DIR', default='./data', help='path to dataset')
parser.add_argument('--epochs', default=120, type=int, metavar='N', help='number of total epochs to run')
parser.add_argument('--batch_size',
                    default=50,
                    type=int,
                    metavar='N',
                    help='mini-batch size (default: 256), this is the total '
                    'batch size of all GPUs on the current node when '
                    'using Data Parallel or Distributed Data Parallel')
parser.add_argument('--lr',
                    '--learning_rate',
                    default=0.05,
                    type=float,
                    metavar='LR',
                    help='initial learning rate',
                    dest='lr')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M', help='Momentum (default: 0.9)', dest='momentum')
parser.add_argument('--weight_decay',
                    default=1e-4,
                    type=float,
                    metavar='W',
                    help='weight decay (default: 1e-4)',
                    dest='weight_decay')
parser.add_argument('--in_dim', default=4096, type=int, help='feature dimension (default: 4096)')
parser.add_argument('--out_dim', default=512, type=int, help='feature dimension (default: 512)')
parser.add_argument('--momentum_update',
                    default=0.996,
                    type=float,
                    metavar='MU',
                    help='Momentum Update (default: 0.996)',
                    dest='momentum_update')
parser.add_argument('--save_model', default=False, type=bool, help='To save model, True or False')
parser.add_argument('--gpus', default=0, type=int, help='Number of Gpus used during training (default: 0)')


def main():
    classes = 48
    args = parser.parse_args()
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    online = ResNetBYOL(projection_in_dimension=args.in_dim,
                        projection_out_dimension=args.out_dim,
                        num_class=classes,
                        mode='online').to(device)
    target = ResNetBYOL(projection_in_dimension=args.in_dim,
                        projection_out_dimension=args.out_dim,
                        num_class=classes,
                        mode='taeget').to(device)
    predictor = MLPHead(input_channel=online.projection[-1].out_features,
                        projection_in_dimension=args.in_dim // 2,
                        projection_out_dimension=args.out_dim).to(device)
    fine_tuning_model = list(online.encoder[1:].parameters())
    custom_model = list(online.encoder[0].parameters()) + list(online.cls.parameters())
    byol_model = list(online.projection.parameters()) + list(predictor.parameters())
    optimizer = torch.optim.SGD([{
        "params": fine_tuning_model
    }, {
        "params": custom_model,
        "lr": args.lr
    }, {
        'params': byol_model,
        'lr': args.lr / 10
    }],
                                lr=args.lr / 5,
                                weight_decay=args.weight_decay,
                                momentum=args.momentum)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)
    if args.gpus > 0:
        online = torch.nn.DataParallel(online, device_ids=range(args.gpus))
        predictor = torch.nn.DataParallel(predictor, device_ids=range(args.gpus))
    trainer = BYOLTrainer(online=online,
                          target=target,
                          predictor=predictor,
                          optimizer=optimizer,
                          scheduler=scheduler,
                          device=device,
                          params=args,
                          classes=classes)
    trainer.train(batch_size=args.batch_size, num_workers=2)
    trainer.save_features('resnet-50')


if __name__ == '__main__':
    main()